{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# PaddlePaddle BYOS\n",
    "\n",
    "## Pre-requisites\n",
    "\n",
    "This notebook shows how to use the SageMaker Python SDK to run your code in a local container before deploying to SageMaker's managed training or hosting environments.  This can speed up iterative testing and debugging while using the same familiar Python SDK interface.  Just change your estimator's `train_instance_type` to `local` (or `local_gpu` if you're using an ml.p2 or ml.p3 notebook instance).\n",
    "\n",
    "In order to use this feature you'll need to install docker-compose (and nvidia-docker if training with a GPU).\n",
    "\n",
    "**Note, you can only run a single local notebook at one time.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !/bin/bash ./utils/setup.sh"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Overview\n",
    "\n",
    "The **SageMaker Python SDK** helps you deploy your models for training and hosting in optimized, productions ready containers in SageMaker. The SageMaker Python SDK is easy to use, modular, extensible and compatible with TensorFlow, MXNet, PyTorch and Chainer. This tutorial focuses on how to create a convolutional neural network model to train the [Cifar10 dataset](https://www.cs.toronto.edu/~kriz/cifar.html) using **PyTorch in local mode**.\n",
    "\n",
    "### Set up the environment\n",
    "\n",
    "This notebook was created and tested on a single ml.p2.xlarge notebook instance.\n",
    "\n",
    "Let's start by specifying:\n",
    "\n",
    "- The S3 bucket and prefix that you want to use for training and model data. This should be within the same region as the Notebook Instance, training, and hosting.\n",
    "- The IAM role arn used to give training and hosting access to your data. See the documentation for how to create these. Note, if more than one role is required for notebook instances, training, and/or hosting, please replace the sagemaker.get_execution_role() with appropriate full IAM role arn string(s)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sagemaker\n",
    "\n",
    "sagemaker_session = sagemaker.Session()\n",
    "\n",
    "bucket = sagemaker_session.default_bucket()\n",
    "prefix = 'sagemaker/DEMO-PaddleNLP'\n",
    "\n",
    "role = sagemaker.get_execution_role()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import subprocess\n",
    "\n",
    "# instance_type = 'local'\n",
    "\n",
    "# if subprocess.call('nvidia-smi') == 0:\n",
    "#     ## Set type to GPU if one is present\n",
    "#     instance_type = 'local_gpu'\n",
    "    \n",
    "# print(\"Instance type = \" + instance_type)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Upload the data\n",
    "We use the ```sagemaker.Session.upload_data``` function to upload our datasets to an S3 location. The return value inputs identifies the location -- we will use this later when we start the training job."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# base_dir = 'file:///home/ec2-user/SageMaker/paddlenlp_sagemaker/data/'\n",
    "# inputs = {'training': base_dir}\n",
    "# print(inputs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Script Functions\n",
    "\n",
    "SageMaker invokes the main function defined within your training script for training. When deploying your trained model to an endpoint, the model_fn() is called to determine how to load your trained model. The model_fn() along with a few other functions list below are called to enable predictions on SageMaker.\n",
    "\n",
    "### [Predicting Functions](https://github.com/aws/sagemaker-pytorch-containers/blob/master/src/sagemaker_pytorch_container/serving.py)\n",
    "* model_fn(model_dir) - loads your model.\n",
    "* input_fn(serialized_input_data, content_type) - deserializes predictions to predict_fn.\n",
    "* output_fn(prediction_output, accept) - serializes predictions from predict_fn.\n",
    "* predict_fn(input_data, model) - calls a model on data deserialized in input_fn.\n",
    "\n",
    "The model_fn() is the only function that doesn't have a default implementation and is required by the user for using PyTorch on SageMaker. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create a training job using the sagemaker.PyTorch estimator\n",
    "\n",
    "The `PyTorch` class allows us to run our training function on SageMaker. We need to configure it with our training script, an IAM role, the number of training instances, and the training instance type. For local training with GPU, we could set this to \"local_gpu\".  In this case, `instance_type` was set above based on your whether you're running a GPU instance.\n",
    "\n",
    "After we've constructed our `PyTorch` object, we fit it using the data we uploaded to S3. Even though we're in local mode, using S3 as our data source makes sense because it maintains consistency with how SageMaker's distributed, managed training ingests data.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# from sagemaker.pytorch import PyTorch\n",
    "\n",
    "# # git_config = {'repo': 'https://github.com/PaddlePaddle/PaddleNLP.git', 'branch': 'develop'}\n",
    "\n",
    "# hyperparameters = {'train_path': '/opt/ml/input/data/training/train.txt', \n",
    "#                    'dev_path': '/opt/ml/input/data/training/dev.txt', \n",
    "#                    'save_dir': '/opt/ml/model', \n",
    "#                    'learning_rate': 1e-5, \n",
    "#                    'batch_size': 16, \n",
    "#                    'max_seq_len':512, \n",
    "#                    'num_epochs': 100, \n",
    "#                    'model': 'uie-base',\n",
    "#                    'seed': 1000,\n",
    "#                    'logging_steps': 10,\n",
    "#                    'valid_steps': 100,\n",
    "#                    'device': 'gpu'}\n",
    "\n",
    "# estimator = PyTorch(entry_point='finetune.py',\n",
    "#                             source_dir='./',\n",
    "# #                             source_dir='model_zoo/uie/',\n",
    "# #                             git_config=git_config,\n",
    "#                             role=role,\n",
    "#                             hyperparameters=hyperparameters,\n",
    "#                             framework_version='1.9.1',\n",
    "#                             py_version='py38',\n",
    "#                             script_mode=True,\n",
    "#                             instance_count=1,  # 1 or 2 or ...\n",
    "#                             instance_type=instance_type)\n",
    "\n",
    "# estimator.fit(inputs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## SageMaker Training using GPU instance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "WORK_DIRECTORY = '/home/ec2-user/SageMaker/paddlenlp_sagemaker/data/'\n",
    "\n",
    "data_location = sagemaker_session.upload_data(WORK_DIRECTORY, key_prefix=prefix)\n",
    "\n",
    "inputs = {'training': data_location}\n",
    "\n",
    "print(inputs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from sagemaker.pytorch import PyTorch\n",
    "\n",
    "hyperparameters = {'train_path': '/opt/ml/input/data/training/train.txt', \n",
    "                   'dev_path': '/opt/ml/input/data/training/dev.txt', \n",
    "                   'save_dir': '/opt/ml/model', \n",
    "                   'learning_rate': 1e-5, \n",
    "                   'batch_size': 16, \n",
    "                   'max_seq_len':512, \n",
    "                   'num_epochs': 100, \n",
    "                   'model': 'uie-base',\n",
    "                   'seed': 1000,\n",
    "                   'logging_steps': 10,\n",
    "                   'valid_steps': 100,\n",
    "                   'device': 'gpu'}\n",
    "\n",
    "instance_type = 'ml.p3.2xlarge'  # 'ml.p3.2xlarge' or 'ml.p3.8xlarge' or ...\n",
    "\n",
    "git_config = {'repo': 'https://github.com/whn09/paddlenlp_sagemaker.git', 'branch': 'main'}\n",
    "\n",
    "estimator = PyTorch(entry_point='finetune.py',\n",
    "                    source_dir='./',\n",
    "#                             source_dir='model_zoo/uie/',\n",
    "                            git_config=git_config,\n",
    "                            role=role,\n",
    "                            hyperparameters=hyperparameters,\n",
    "                            framework_version='1.9.1',\n",
    "                            py_version='py38',\n",
    "                            script_mode=True,\n",
    "                            instance_count=1,  # 1 or 2 or ...\n",
    "                            instance_type=instance_type)\n",
    "\n",
    "estimator.fit(inputs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "training_job_name = estimator.latest_training_job.name\n",
    "# training_job_name = 'xxx'\n",
    "print(training_job_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Deploy the trained model to prepare for predictions\n",
    "\n",
    "The deploy() method creates an endpoint (in this case locally) which serves prediction requests in real-time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!rm -rf model.tar.gz\n",
    "!rm -rf model_*\n",
    "!rm -rf inference.*\n",
    "!aws s3 cp s3://$bucket/$training_job_name/output/model.tar.gz .\n",
    "!tar -xvf model.tar.gz"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!cp inference.* model/\n",
    "!cp model/code/requirements_gpu.txt model/code/requirements.txt\n",
    "!cd model && tar -czvf ../model-inference-gpu.tar.gz *\n",
    "\n",
    "!aws s3 cp model-inference-gpu.tar.gz s3://$bucket/$training_job_name/output/model-inference-gpu.tar.gz"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# instance_type = 'local'\n",
    "# instance_type = 'ml.m5.xlarge'\n",
    "instance_type = 'ml.g4dn.xlarge'\n",
    "\n",
    "# predictor = estimator.deploy(initial_instance_count=1, instance_type=instance_type)\n",
    "\n",
    "from sagemaker.pytorch.model import PyTorchModel\n",
    "\n",
    "pytorch_model = PyTorchModel(model_data='s3://{}/{}/output/model-inference-gpu.tar.gz'.format(bucket, training_job_name), role=role,\n",
    "                             entry_point='infer_gpu.py', framework_version='1.9.0', py_version='py38', model_server_workers=4)  # TODO [For GPU], model_server_workers=6\n",
    "\n",
    "predictor = pytorch_model.deploy(instance_type=instance_type, initial_instance_count=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # endpoint_name = 'pytorch-inference-2022-07-05-07-28-16-183'  # m5.2xlarge\n",
    "# # endpoint_name = 'pytorch-inference-2022-07-06-04-02-11-091'  # g4dn.xlarge, 6 threads\n",
    "# endpoint_name = 'pytorch-inference-2022-07-06-06-19-21-855'  # g4dn.xlarge, 4 threads\n",
    "# predictor = sagemaker.predictor.Predictor(endpoint_name=endpoint_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Invoking the endpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker.serializers import JSONSerializer\n",
    "from sagemaker.deserializers import JSONDeserializer\n",
    "\n",
    "predictor.serializer = JSONSerializer()\n",
    "predictor.deserializer = JSONDeserializer()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "texts = ['\"北京市海淀区人民法院\\n民事判决书\\n(199x)建初字第xxx号\\n原告：张三。\\n委托代理人李四，北京市 A律师事务所律师。\\n被告：B公司，法定代表人王五，开发公司总经理。\\n委托代理人赵六，北京市 C律师事务所律师。\"', \n",
    "         '原告赵六，2022年5月29日生\\n委托代理人孙七，深圳市C律师事务所律师。\\n被告周八，1990年7月28日出生\\n委托代理人吴九，山东D律师事务所律师']\n",
    "\n",
    "import time\n",
    "start = time.time()\n",
    "outputs = predictor.predict(texts)\n",
    "end = time.time()\n",
    "print('outputs: ', outputs)\n",
    "print('time:', end-start)\n",
    "\n",
    "# for i in range(1000):\n",
    "#     start = time.time()\n",
    "#     outputs = predictor.predict(texts)\n",
    "#     end = time.time()\n",
    "#     print('time:', end-start)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Clean-up\n",
    "\n",
    "Deleting the local endpoint when you're finished is important since you can only run one local endpoint at a time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# estimator.delete_endpoint()\n",
    "predictor.delete_endpoint()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "conda_python3",
   "language": "python",
   "name": "conda_python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.13"
  },
  "notice": "Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.  Licensed under the Apache License, Version 2.0 (the \"License\"). You may not use this file except in compliance with the License. A copy of the License is located at http://aws.amazon.com/apache2.0/ or in the \"license\" file accompanying this file. This file is distributed on an \"AS IS\" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License."
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
